#include "Dwarf.Unmanaged.Mpi.h"

Solver :: Solver()
{   
    MPI_Comm_rank(MPI_COMM_WORLD, &commRank);           // Retrieve rank of the process
    MPI_Comm_size(MPI_COMM_WORLD, &commSize);           // Retrieve count of processes
    arrayOfProccLengths = new int[commSize];
    arrayOfShifts = new int[commSize];
    isRootThread = commRank == NUMBER_ROOT_PROCESS;
}

Solver :: ~Solver()
{    
    delete hmm;
    delete vit;
    delete arrayOfProccLengths;
    delete arrayOfShifts;
}

void Solver :: init(int n, int m, int t)
{
    this->hmm = new HMM(n,m);                   //creation of Hidden Markov model  instance
    this->vit = new ViterbiModel(t,this->hmm);  //creation of Viterbi Model instance
}

//For the demonstration purposes of the project,
//we have chosen Viterbi algorithm for Hidden Markov Model
void Solver::solve() 
{
    int maxvalind;                         //temporary storage for max value for psi array
    double maxval, val;                    //temporary storage for max value for delta array
    int numberOfStates = hmm->getNumberOfStates();

    if(isRootThread)
    {
        /* 1. Initialization delta and psi */
        for (int i = 0; i < numberOfStates; i++) 
        {
            vit->delta[0][i] = hmm->initialStates[i] * (hmm->observationProbabilityMatrix[i][vit->observationSequence[0]]);
            vit->psi[0][i] = 0;
        }
    }
    /* 2. Compute*/

    //sending to all instances the initial delta and psi values
    MPI_Bcast(vit->delta[0], numberOfStates, MPI_DOUBLE, NUMBER_ROOT_PROCESS , MPI_COMM_WORLD);

    int lengthOfObsSeq = vit->getLengthOfObservationSequence();
    double * workDeltaArray =  new double[lengthForProc];       //temporary storage of delta values
    int * workPsiArray = new int[lengthForProc];                //temporary storage of psi values
    int p;                                                      //helper index
    for (int t = 1; t < lengthOfObsSeq; t++)
    {   
        //wait for all proc
        MPI_Barrier(MPI_COMM_WORLD);
        p = 0;
        for (int j = startIndx; j < endIndx; j++) 
        {
            maxval = 0.0;
            maxvalind = 0;
            //finding the maximum
            for (int i = 0; i < numberOfStates; i++) 
            {
                val = vit->delta[t-1][i]*(hmm->stateTransitionMatrix[i][j]);
                if (val > maxval) 
                {
                    maxval = val;
                    maxvalind = i;
                }
            }

            workDeltaArray[p] = maxval*(hmm->observationProbabilityMatrix[j][vit->observationSequence[t]]);
            workPsiArray[p] = maxvalind;
            p++;
        }
        //send to all instances the results of computation on step t
        MPI_Allgatherv(workDeltaArray, lengthForProc, MPI_DOUBLE, vit->delta[t], arrayOfProccLengths, arrayOfShifts, MPI_DOUBLE, MPI_COMM_WORLD);
        MPI_Gatherv(workPsiArray, lengthForProc, MPI_INT, vit->psi[t], arrayOfProccLengths, arrayOfShifts, MPI_INT, NUMBER_ROOT_PROCESS, MPI_COMM_WORLD);
    }
    //final computation on master process
    if(isRootThread)        
    {   
        /* 3. Termination */
        vit->setProbability(0);

        vit->sequenceOfState[lengthOfObsSeq - 1] = 0;
        for (int i = 0; i < numberOfStates; i++) 
        {   
            if (vit->delta[lengthOfObsSeq - 1][i] > vit->getProbability()) 
            {   
                vit->setProbability(vit->delta[lengthOfObsSeq - 1][i]);
                vit->sequenceOfState[lengthOfObsSeq - 1] = i;
            }
        }

        /* 4. Path (state sequence) backtracking */
        for (int t = lengthOfObsSeq - 2; t >= 0; t--)
        {
            vit->sequenceOfState[t] = vit->psi[t+1][vit->sequenceOfState[t+1]];
        }
    } 

    delete workDeltaArray;
    delete workPsiArray;
}